from langchain.chains import RetrievalQA 
from vectores import GeneratedVectors
import time
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_core.output_parsers import StrOutputParser

DASHSCOPE_API_KEY = "sk-2c282b84a1b34473a136e3e34229d062"
time_start = time.time()
import os
from langchain.vectorstores import FAISS
from langchain_community.llms import Tongyi
# 初始化 DashScope 大语言模型（用于问答）
llm = Tongyi(model_name="deepseek-v3", dashscope_api_key=DASHSCOPE_API_KEY)

# 创建检索器（假设您已经创建了 FAISS 向量数据库）
GenerateVector = GeneratedVectors()
vectordb = GenerateVector.get_vectors()
from langchain.retrievers import MultiQueryRetriever

retriever = MultiQueryRetriever.from_llm(
    retriever = vectordb.as_retriever(),
    llm=llm
)

# 
from langchain.prompts import PromptTemplate
from langchain.chains import RetrievalQA
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'  # 注意要写在最开头

#对话检索链
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate

class LangchainQA:
    def __init__(self):
        # 将llm和retriever保存为实例变量
        self.llm = llm
        self.retriever = retriever
        
        # 自定义提示模板
        self.prompt_template = """你是一名高校合格评估问答AI助手，你可以基于以下上下文信息和历史信息来回答问题。如果你不知道答案，就说不知道，不要编造答案。
        对话历史：
        {chat_history}

        上下文:
        {context}

        问题: {question}
        答案:"""

        self.PROMPT = PromptTemplate(
            template=self.prompt_template, 
            input_variables=["chat_history", "context", "question"]
        )
        self.memory = ConversationBufferMemory(
            memory_key="chat_history", 
            return_messages=True,
            output_key="answer"
        )

    # 创建带有自定义提示的检索式问答链
    def create_chain(self):
        qa = ConversationalRetrievalChain.from_llm(
            llm=self.llm,  # 使用实例变量
            memory=self.memory,  # 使用实例变量默认的内存
            retriever=self.retriever,  # 使用实例变量
            combine_docs_chain_kwargs={"prompt": self.PROMPT}, # 
            return_source_documents=True,
            verbose=True
        )
        return qa


class custom_rag_chain:
    def __init__(self):
         # 将llm和retriever保存为实例变量
        self.llm = llm
        self.retriever = retriever
        
        # 自定义提示模板
        self.prompt_template = """基于以下上下文信息，请回答问题。如果你不知道答案，就说不知道，不要编造答案。

上下文:
{context}

问题: {question}
答案:"""
        self.PROMPT = PromptTemplate(
            template=self.prompt_template, 
            input_variables=["context", "question"]

        )
        #优化查询
        from opt_retriveral import GeneratedQueryDoc
        self.opt_retriveral = GeneratedQueryDoc()
        # 1. 创建提示模板
        # 2. 定义一个新的RAG链，在检索前先进行查询扩展
        # 查询扩展方法
        # 2. 创建标准的 RetrievalQA 链
        self.qa_chain = RetrievalQA.from_chain_type(
        llm=llm, # 您的语言模型
        chain_type="stuff",
        retriever=retriever, # 您的检索器
        chain_type_kwargs={"prompt": self.PROMPT}, # 注入自定义提示
        return_source_documents=True, # 返回源文档
        verbose=True
    )
    # 3. 创建一个外部的“包装器”函数来处理查询扩展
    def query2doc_qa_chain(self, query):
        """
        包装函数：先进行查询扩展，再调用QA链
        """
        # 3.1 进行查询扩展 (您的Query2Doc逻辑)
        query2doc_chain = self.opt_retriveral.get_retriever()
        expanded_query = query2doc_chain.invoke({"query": query})
        return expanded_query
    def get_result(self, query):
        # 3.2 使用扩展后的查询进行检索和问答
        expanded_query = self.query2doc_qa_chain(query)
        # print("expended_query:", expanded_query)
        result = self.qa_chain({"query": expanded_query})
        # print("result:", result)
        return result

        
#写一个测试
if __name__ == "__main__":
    langchain_qa = LangchainQA()
    qa = langchain_qa.create_chain()
    # new_langchain = custom_rag_chain()
    query = "合格评估的依据？"
    start_time = time.time()
    # ans = new_langchain.get_result()
    # print("qa", qa)
    ans = qa.invoke({"question": query})
    # print("ans",ans)
    end_time = time.time()
    total_time = f"{round(end_time - start_time, 3)} sec."
    # 提取答案和源文档
    response = ans['answer']
    source_docs = ans.get('source_documents', [])
    
    # 构建完整响应
    full_response = f"Question: {query}\nAnswer: {response}\nTotal time: {total_time}"
    
    # 添加源文档信息（如果有）
    if source_docs:
        full_response += "\n\n来源文档:\n"
        for i, doc in enumerate(source_docs):
            source_info = f"{doc.metadata.get('source', '未知')}, 页码: {doc.metadata.get('page', '未知')}"
            full_response += f"{i+1}. {doc.page_content[:100]}... [{source_info}]\n"
    
    print(full_response)
